﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using Microsoft.Win32.SafeHandles;
using Mono.Cecil.PE;
using NHook.Native;

namespace NHook
{
	public enum Instr
	{
		ShortJMP = 0xEB,
		LongJMP = 0xE9,
		Mov = 0x8B,
		MovECX = 0xB9,
		Ecx = 0x8D,
		Lea = 0x8D,
		JMPF = 0xEA,
		PushECX = 0x51,
		PushESI = 0x56,
		Nop = 0x90
	}

	public class Patch : IDisposable
	{
		private Stream reader;
		private Stream writer;
		private MemoryStream copy = new MemoryStream();
		IntPtr start = IntPtr.Zero;
		bool _Undid = false;


		public Patch(SafeFileHandle processHandle, IntPtr addr)
		{
			CurrentAddress = addr;
			start = addr;
			this.reader = new RemoteUnmanagedMemoryStream(processHandle, addr);
			this.writer = new RemoteUnmanagedMemoryStream(processHandle, addr);
		}

		public unsafe Patch(IntPtr addr)
		{
			CurrentAddress = addr;
			start = addr;
			this.reader = new UnmanagedMemoryStream((byte*)addr.ToPointer(), Int32.MaxValue);
			this.writer = new UnmanagedMemoryStream((byte*)addr.ToPointer(), Int32.MaxValue, Int32.MaxValue, FileAccess.Write);
		}

		public void Mov(Instr instruction, IntPtr source)
		{
			WriteByte((byte)instruction);
			//WriteByte((byte)target);
			WriteReverse((uint)source.ToInt32());
		}

		public void Jump(Instr instruction, IntPtr from, IntPtr to)
		{
			WriteByte((byte)instruction);

			uint jmpAddr = 0;
			if(to.ToInt32() < from.ToInt32())
				jmpAddr = unchecked((uint)(to.ToInt32() - from.ToInt32())) - 5;
			else
				jmpAddr = unchecked((uint)(to.ToInt32() - from.ToInt32())) + 5;
			//jmpAddr = addr.ToInt32();
			WriteReverse(jmpAddr);
		}

		private void WriteReverse(uint addr)
		{
			for(int i = 0 ; i < 4 ; i++)
			{
				var byteAddr = (byte)((addr >> i * 8) & 0xFF);
				WriteByte(byteAddr);
			}
		}

		private void WriteByte(byte value)
		{
			copy.WriteByte((byte)reader.ReadByte());
			writer.WriteByte((byte)value);
			CurrentAddress = new IntPtr(CurrentAddress.ToInt32() + 1);
			writer.Flush();
		}
		private void WriteBytes(byte[] buffer)
		{
			for(int i = 0 ; i < buffer.Length ; i++)
				WriteByte(buffer[i]);
		}

		public DisasmResult Write(string instruction)
		{
			var result = DisasmWrapper.Assemble(instruction);
			WriteBytes(result.Bytes);
			return result;
		}

		#region IDisposable Members

		public void Dispose()
		{
			Undo();
		}

		#endregion

		public void Undo()
		{
			if(_Undid)
				return;
			PatchFromCopy();
			_Undid = true;
		}

		public void Redo()
		{
			if(!_Undid)
				return;
			PatchFromCopy();
			_Undid = false;
		}

		private unsafe void PatchFromCopy()
		{
			copy.Position = 0;
			int value = 0;
			writer.Position = start.ToInt64();
			reader.Position = start.ToInt64();
			do
			{
				value = copy.ReadByte();
				if(value == -1)
					break;
				copy.Position--;
				copy.WriteByte((byte)reader.ReadByte());
				writer.WriteByte((byte)value);
			} while(value != -1);
		}

		public void Write(Instr instr)
		{
			WriteByte((byte)instr);
		}

		[DllImport("kernel32.dll")]
		public static extern bool FlushInstructionCache(IntPtr hProcess, IntPtr lpBaseAddress, UIntPtr dwSize);


		public void Flush()
		{
			FlushInstructionCache(System.Diagnostics.Process.GetCurrentProcess().Handle, start, (UIntPtr)(CurrentAddress.ToInt32() - start.ToInt32()));
		}

		public IntPtr CurrentAddress
		{
			get;
			private set;
		}

		public bool ApplyPatch(ProcessModule module, Stream stream)
		{
			var image = Image.ReadFromFile(module.FileName);
			var rvaInt = start.ToInt32() - module.BaseAddress.ToInt32();
			if(rvaInt <= 0)
				return false;
			var rva = new RVA((uint)rvaInt);
			var offset = image.ToImageOffset(rva);
			stream.Position = offset.Address;

			var readOldPosition = reader.Position;
			reader.Position = start.ToInt64();
			try
			{
				while(readOldPosition != reader.Position)
					stream.WriteByte((byte)reader.ReadByte());
			}
			finally
			{
				reader.Position = readOldPosition;
			}
			stream.Flush();
			return true;
		}

		
	}
}
